# -*- coding:utf-8 -*-
# created on 2017/01/05

from mathsolver.functions.base import *
import numpy as np
from sympy import Symbol, symbols, cos, acos, Abs, Mul, Add, Pow, S, pi, sqrt, Integer


# 求向量夹角
class XiangLiangMoCompute001(BaseFunction):
    def solver(self, *args):
        vector_a = sympify(args[0].value)  # 已知的两个向量名称
        vector_b = sympify(args[1].value)
        abs_ = self.search('Abs')  # 寻找所有的向量模
        arith = self.search('Arith')  # 寻找所有的向量四则运算
        theta, angle = symbols('theta angle')
        # 有些情况向量的夹角已经通过其他函数运算保存到了内存中，所以如果匹配到就可以直接使用了
        if self.search(angle):
            self.output.append(BaseValue(self.search(angle)))
        else:
            self.steps.append(["", "设%s和%s的夹角是%s" % (str(vector_a), str(vector_b), str(theta))])
            self.steps.append(["", "根据已知内容可计算得以下方程组:"])
            eqs = []  # 将所有向量预算转换为方程，最后变成解方程组求theta
            flag = 0
            if self.search(vector_a):
                a = self.search(vector_a)
                if isinstance(a, int):
                    a_value = np.array(self.search(a))
                else:
                    flag += 1
                    a_value = sqrt(a[0] ** 2 + a[1] ** 2)
                eqs.append(vector_a - a_value)
            if self.search(vector_b):
                b = self.search(vector_b)
                if isinstance(a, int):
                    b_value = np.array(self.search(vector_b))
                else:
                    flag += 1
                    b_value = sqrt(b[0] ** 2 + b[1] ** 2)
                eqs.append(vector_b - b_value)
            if flag == 2:
                e = vector_a * vector_b - sum(np.array(a) * np.array(b))
                eqs.append(e)
            if abs_:
                for s in abs_:
                    eqs.append(s[0].args[0] - s[1])
            if arith:
                for h in arith:
                    e = h[0] - h[1]
                    e = e.expand()
                    e_new = S.Zero
                    co = e.as_coefficients_dict()
                    # 将所有的变量替换为n*vector(a)*vector(b)的形式，需要把a和b的系数提出来
                    for cc in co:
                        if isinstance(cc, Mul) or isinstance(cc, Pow):  # 如果该变量是vector(a)*vector(b)的格式
                            for c in cc.args:  # 读取每一个vector(a)和vector(b)
                                if str(c.func == 'vector'):
                                    c_new = XiangLiangMoCompute002().vector_expand(c)
                                    e_new = e.subs(c, c_new)
                        else:  # 如果是一个常数
                            e_new = e_new + cc
                    eqs.append(e_new)

            # 将方程组中的a*b替换为a*b*theta，因为向量点积需要考虑夹角
            for i in range(len(eqs)):
                eqs[i] = eqs[i].subs(vector_a * vector_b, vector_a * vector_b * theta)
                self.steps.append(["", "%s = 0" % str(eqs[i])])
            cores_set = solve(eqs, [vector_a, vector_b, theta], set=True)
            cores_list = list(cores_set[1])
            cores = cores_list[0][2]
            res = acos(cores)

            self.output.append(BaseValue(res))

        return self


# 根据向量模长和夹角求未知向量
class XiangLiangMoCompute002(BaseFunction):
    # 将vector(2*a)变为2*vector(a)
    @staticmethod
    def vector_expand(arg):
        vector = arg
        if isinstance(vector, Mul):
            vector_new = vector.args[0] * vector.args[1]
        else:
            vector_new = vector
        return vector_new

    def solver(self, *args):
        vector = args[0].value
        vector_a = sympify(vector[0])
        vector_b = sympify(vector[1])
        vectoreqs = sympify(args[1].value)
        theta, angle = symbols('theta ,angle')
        eqs = []
        # 判断向量等式组中的参数类型
        for veq in vectoreqs:
            if isinstance(veq[0], Abs):  # 如果参数是模长
                if len(veq) == 3:  # 两个向量的模长一样
                    e1 = veq[0].args[0] - veq[2]
                    e2 = veq[1].args[0] - veq[2]
                    eqs.append(e1)
                    eqs.append(e2)
                else:
                    if str(veq[0].args[0].func) == 'vector':  # 一个向量的模长a
                        e = veq[0].args[0] - veq[1]
                        eqs.append(e)
                    else:  # 向量运算后的模长a+b
                        e = veq[0].args[0] ** 2 - veq[1] ** 2
                        e = e.expand()
                        eqs.append(e)
            elif str(veq[0].func) == 'vector':  # 如果参数是向量
                e = veq[0] - veq[1]
                eqs.append(e)
            elif isinstance(veq[0], Mul) or isinstance(veq[0], Add):  # 如果参数是向量的四则运算
                e = veq[0] - veq[1] if isinstance(veq[1], int) else veq[0] - 0
                e = e.expand()
                e_new = 0
                co = e.as_coefficients_dict()
                # 将所有的变量替换为n*vector(a)*vector(b)的形式，需要把a和b的系数提出来
                for cc in co:
                    if isinstance(cc, Mul) or isinstance(cc, Pow):  # 如果该变量是vector(a)*vector(b)的格式
                        for c in cc.args:  # 读取每一个vector(a)和vector(b)
                            if str(c.func == 'vector'):
                                c_new = XiangLiangMoCompute002().vector_expand(c)
                                e_new = e.subs(c, c_new)
                    else:  # 如果是一个常数
                        e_new = e_new + cc
                eqs.append(e_new)
            else:  # 如果参数是夹角
                theta = veq[1]
                if isinstance(theta, Integer):  # theta是度数的形式
                    xldj = vector_a * vector_b * cos(pi * theta / 180)
                else:  # theta是带pi的形式
                    xldj = vector_a * vector_b * cos(theta)
        self.steps.append(["", "根据题意可得到方程组:"])
        # 将方程组中的a*b替换为a*b*theta，因为向量点积需要考虑夹角
        for i in range(len(eqs)):
            if isinstance(theta, Symbol):
                eqs[i] = eqs[i].subs(vector_a * vector_b, vector_a * vector_b * theta)
            else:
                eqs[i] = eqs[i].subs(vector_a * vector_b, xldj)
            self.steps.append(["", "%s = 0" % str(eqs[i])])

        dict_items = {}
        if isinstance(theta, Symbol):
            res_set = solve(eqs, [vector_a, vector_b, theta], set=True)
            res = list(res_set[1])
            for r in res:
                if r[0] * r[1] >= 0:
                    dict_items[vector_a] = r[0]
                    dict_items[vector_b] = r[1]
                    dict_items[angle] = [vector_a, vector_b, acos(r[2])]
        else:
            res_set = solve(eqs, [vector_a, vector_b], set=True)
            res = list(res_set[1])
            for r in res:
                if r[0] * r[1] > 0:
                    dict_items[vector_a] = r[0]
                    dict_items[vector_b] = r[1]
            dict_items[angle] = [vector_a, vector_b, theta]
        self.output.append(BaseSymbolValue(dict_items))
        return self
